# -*- coding: utf-8 -*-
# @version : 1.0
# @File    : auth.py
# @Date    : 2024/9/27 23:56
# @Author  : GlenZhao
# @Desc    :
import json
from typing import Annotated
from fastapi import Request, Depends, WebSocket
from fastapi.exceptions import HTTPException
from fastapi.security.utils import get_authorization_scheme_param
from starlette.status import HTTP_401_UNAUTHORIZED
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import joinedload

from application import settings, constant
from apps.auth.user.models import AuthUser
from apps.auth.user.utils.login_manage import LoginManage
from apps.auth.user.utils.validation import AuthValidation, Auth, WebsocketAuth
from apps.auth.user.crud import UserDal
from apps.auth.user.schemas import UserCache
from apps.auth.user.params import UserAuthParams
from core.database import db_getter, redis_getter
from core.redis import RedisManager


class UserAuth(AuthValidation):
    """
    支持所有用户认证
    获取用户基本信息
    """
    with_user = False # 取用户信息，不走缓存
    with_role = False # 取用户角色

    def __init__(
            self,
            with_user: bool | None = False,
            with_role: bool | None = False
    ):
        self.with_user = with_user
        self.with_role = with_role
        if with_role:
            self.with_user = True

    async def __call__(
            self,
            request: Request,
            token: Annotated[str, Depends(settings.oauth2_scheme)],
            db: AsyncSession = Depends(db_getter),
            redis_manager: RedisManager = Depends(redis_getter)
    ):
        if not settings.OAUTH_ENABLE:
            return Auth(db=db)
        id = self.validate_token(request, token)

        user = None
        user_cache = None

        if not self.with_user:
            user_cache_json = await redis_manager.get(LoginManage.gen_user_key(id))
            if user_cache_json:
                user_cache = UserCache(**json.loads(user_cache_json))

        if self.with_user or not user_cache:
            v_options = []
            if self.with_role:
                v_options.append(joinedload(AuthUser.roles))
            if len(v_options) == 0:
                v_options = None
            user = await UserDal(db).get_data(id, v_options=v_options, v_return_none=True)
            if user and self.with_role:
                user_cache = await LoginManage.cache_userinfo(user)

        if self.with_user and not user:
            AuthValidation.forbidden("用户不存在")

        return await self.validate_user(request, user_cache, user, db, token)


class UserQueryAuth(AuthValidation):
    async def __call__(
        self,
        params: UserAuthParams = Depends(),
        db: AsyncSession = Depends(db_getter),
        redis_manager: RedisManager = Depends(redis_getter)
    ):
        id = self.validate_token(request=None, token=params.token)

        user = None
        user_cache = None

        user_cache_json = await redis_manager.get(LoginManage.gen_user_key(id))
        if user_cache_json:
            user_cache = UserCache(**json.loads(user_cache_json))
        if not user_cache:
            user = await UserDal(db).get_data(id, v_return_none=True)

        if not user and not user_cache:
            AuthValidation.forbidden("用户不存在")
        return Auth(user=user, user_cache=user_cache, db=db, headers=None, token=params.token) 


class UserSocketAuth(AuthValidation):

    async def __call__(
        self,
        websocket: WebSocket,
        db: AsyncSession = Depends(db_getter)
    ):
        headers = websocket.headers
        authorization = headers.get("authorization")
        scheme, access_token = get_authorization_scheme_param(authorization)
        if not authorization or scheme.lower() != "bearer":
            query_params = websocket.query_params
            access_token = query_params.get("access_token")

        if not access_token:
            raise HTTPException(
                status_code=HTTP_401_UNAUTHORIZED,
                detail="Not authenticated",
                headers={"WWW-Authenticate": "Bearer"},
            )
        id = self.validate_token(request=None, token=access_token)
        user = await UserDal(db).get_data(id, v_return_none=True)

        if not user:
            AuthValidation.forbidden("用户不存在")

        return WebsocketAuth(websocket=websocket, db=db, user=user)


